import cv2
from ultralytics import YOLO

# 加载预训练的 YOLOv8 模型（默认是 yolov8n，轻量但速度快）
model = YOLO("yolov8n.pt")  # 可换为 yolov8s.pt / yolov8m.pt 精度更高

# 打开视频
cap = cv2.VideoCapture("video1.mp4")
if not cap.isOpened():
    print("无法打开视频")
    exit()

# 获取帧率、尺寸
fps = cap.get(cv2.CAP_PROP_FPS)
w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))

# 输出视频（带检测框）
out = cv2.VideoWriter("video1_detected.mp4",
                      cv2.VideoWriter_fourcc(*'mp4v'),
                      fps, (w, h))

frame_id = 0
while True:
    ret, frame = cap.read()
    if not ret:
        break
    frame_id += 1

    # YOLOv8 检测（返回结果列表）
    results = model(frame)

    # 可视化人体框
    for r in results:
        for box in r.boxes:
            cls_id = int(box.cls[0])
            conf = float(box.conf[0])
            if cls_id == 0 and conf > 0.5:  # cls 0 表示 "person"
                x1, y1, x2, y2 = map(int, box.xyxy[0])
                cx, cy = (x1 + x2) // 2, (y1 + y2) // 2  # 中心点
                cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 255, 0), 2)
                cv2.circle(frame, (cx, cy), 4, (0, 0, 255), -1)
                cv2.putText(frame, f"Person ({conf:.2f})", (x1, y1 - 10),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)

    # 显示和写入输出视频
    cv2.imshow("Detection", frame)
    out.write(frame)

    # 按 'q' 退出
    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
out.release()
cv2.destroyAllWindows()
